package com.boge.ai.utils;

import com.boge.ai.entity.ChatCompletionResponse;
import okhttp3.*;
import org.apache.commons.lang3.StringEscapeUtils;

import java.util.concurrent.TimeUnit;

/**
 * 大模型操作的公共方法
 */
public class LLMUtils {

    private static final String BASE_URL = "https://api.openai-hk.com/v1/chat/completions";

    private static final String API_KEY = "hk-w3q3id1000054953c829f33c861457133e3048fa220f1076";

    public static String completion(String prompt,String model){
        OkHttpClient client = new OkHttpClient().newBuilder()
                .connectTimeout(20, TimeUnit.SECONDS)
                .readTimeout(20, TimeUnit.SECONDS)
                .build();
        MediaType mediaType = MediaType.parse("application/json");

        String messages = String.format("""
                {
                    "messages":
                        [
                            {
                             "role": "user",
                             "content": "%s"
                             }],
                    "model":"%s"
                }
                """, StringEscapeUtils.escapeJson(prompt), model);
        //System.out.println(messages);
        RequestBody body = RequestBody.create(mediaType, messages);
        //System.out.println(body.toString());
        Request request = new Request.Builder()
                .url(BASE_URL)
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .addHeader("Authorization", "Bearer "+API_KEY)
                .build();
        Response response = null;
        try {
            response = client.newCall(request).execute();
            ResponseBody responseBody = response.body();
            String jsonString = responseBody.string();
            System.out.println(jsonString);
            ChatCompletionResponse chat = JsonToHashMapUtils.parseJsonToResponse(jsonString);


            // 获取嵌套字段：例如 choices[0].message.content
            return chat.getChoices().get(0).getMessage().getContent();
        }catch (Exception e){
            e.printStackTrace();
        }finally {
            response.close();
        }
        return null;
    }
}
